package fun.ticsmyc.test;

import fun.ticsmyc.tools.ArrayTools;
import org.junit.Test;

import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;

public class lc2503 {

    int[] f;
    int[] score;

    public int[] maxPoints(int[][] grid, int[] queries) {
        f = new int[grid.length * grid[0].length];
        score = new int[f.length];
        Arrays.fill(score, 1);
        for (int i = 0; i < f.length; i++) {
            f[i] = i;
        }

        PriorityQueue<int[]> edges = new PriorityQueue<>(Comparator.comparingInt(a -> a[2]));
        for (int i = 0; i < grid.length; ++i) {
            for (int j = 0; j < grid[0].length; ++j) {
                if (i + 1 < grid.length) {
                    edges.add(new int[]{getIndex(grid, i, j), getIndex(grid, i + 1, j), Math.max(grid[i][j], grid[i + 1][j])});
                }
                if (j + 1 < grid[0].length) {
                    edges.add(new int[]{getIndex(grid, i, j), getIndex(grid, i, j + 1), Math.max(grid[i][j], grid[i][j + 1])});
                }
            }
        }

        Integer[] querySort = new Integer[queries.length];
        for (int i = 0; i < querySort.length; i++) {
            querySort[i] = i;
        }
        Arrays.sort(querySort, Comparator.comparingInt(a -> queries[a]));

        int[] res = new int[queries.length];
        for (int i = 0; i < querySort.length; ++i) {
            int query = queries[querySort[i]];
            while (!edges.isEmpty() && edges.peek()[2] < query) {
                int[] poll = edges.poll();
                union(poll[0], poll[1]);
            }
            if (score[find(0)] == 1) {
                res[querySort[i]] = query < grid[0][0] ? 0 : 1;
            } else {
                res[querySort[i]] = score[find(0)];
            }
        }
        return res;
    }


    private int getIndex(int[][] grid, int x, int y) {
        return grid[0].length * x + y;
    }

    public int find(int a) {
        if (f[a] == a) {
            return a;
        }
        return f[a] = find(f[a]);
    }

    public void union(int a, int b) {
        int i = find(a);
        int i1 = find(b);
        if (i != i1) {
            f[i] = i1;
            int add = score[i1] + score[i];
            score[i1] = add;
            score[i] = add;
        }
    }

    @Test
    public void test() {
        ArrayTools.printArray(maxPoints(new int[][]{{1, 2, 3}, {2, 5, 7}, {3, 5, 1}}, new int[]{5, 6, 2}));
    }
}
